import {Embeddings} from 'langchain/embeddings/base'
import axios from 'axios'

export class MyEmbeddings extends Embeddings {
    constructor(fields, configuration) {
        super(fields ?? {});
    }
    async embedDocuments(texts) {
        const subPrompts = texts.map((t) => t.replace(/(\n)|(\r)|(\f)|( )/g, ""))
        const embeddings = [];
        for (let i = 0; i < subPrompts.length; i += 1) {
            const input = subPrompts[i];
            const { data } = await this.embed( input );
            embeddings.push(data);
        }
        return embeddings;
    }
    async embedQuery(text) {
        const { data } = await this.embed( text );
        return data
    }
    async embed(text, maxRetries = 5, timeout = 5000) {
        let retryCount = 0;
        while (retryCount < maxRetries) {
            try {
                const res = await axios.post('http://43.134.44.205:30626/main_cli', {
                    "type": "passage",
                    "text": text
                }, {
                    timeout: timeout,
                });
                return {
                    data: res.data.para_embs[0]
                };
            } catch (error) {
                retryCount++;
            }
        }
        throw new Error(`达到最大重试次数(${maxRetries}次)`);
    }
}
